import z3
def free_single(x, y):
    square_sum = x**2 + y**2
    ring1 = z3.And(0.5**2< square_sum, square_sum < 1**2)
    broken_ring1 = z3.And(ring1, y<0.5)
    
    ring2 = z3.And(1.5**2< square_sum, square_sum < 2**2)
    broken_ring2 = z3.And(ring2, y>-1.5)
    
    obs = z3.Or(broken_ring1, broken_ring2)
    free = z3.Not(obs)
    return free

def free_single_linear(x, y):
    belt1 = z3.And(x-y > -1,  x-y < -0.5, x<0,  y>0,  y<0.25)
    belt2 = z3.And(x+y > 0.5, x+y < 1,    x>=0, y>0,  y<0.25)
    belt3 = z3.And(x-y > 0.5, x-y < 1,    x>=0, y<=0)
    belt4 = z3.And(x+y > -1,  x+y < -0.5, x<0,  y<=0)
    broken_square1 = z3.Or(belt1, belt2, belt3, belt4)
    
    belt1 = z3.And(x-y > -2,  x-y < -1.5, x<0,  y>0)
    belt2 = z3.And(x+y > 1.5, x+y < 2,    x>=0, y>0)
    belt3 = z3.And(x-y > 1.5, x-y < 2,    x>=0, y<=0, y>-1.25)
    belt4 = z3.And(x+y > -2,  x+y < -1.5, x<0,  y<=0, y>-1.25)
    broken_square2 = z3.Or(belt1, belt2, belt3, belt4)
    
    obs = z3.Or(broken_square1, broken_square2)
    free = z3.Not(obs)
    return free

def free_single_easy(x, y):
    belt1 = z3.And(x < -0.5, x > -1,  y < 1, y > -1)
    belt2 = z3.And(x < 1,    x > 0.5, y < 1, y > -1)
    belt3 = z3.And(y < -0.5, y > -1,  x < 1, x > -1)
    broken_square1 = z3.Or(belt1, belt2, belt3)
    belt1 = z3.And(x < -1.5, x > -2,  y < 2, y > -2)
    belt2 = z3.And(x < 2,    x > 1.5, y < 2, y > -2)
    belt3 = z3.And(y < 2, y > 1.5,  x < 2, x > -2)
    broken_square2 = z3.Or(belt1, belt2, belt3)

    obs = z3.Or(broken_square1, broken_square2)
    free = z3.Not(obs)
    return free

def free_constraint(x, y, t, single_constraint):
    N = len(x)
    c = []
    for i in range(N-1):
        c.append(single_constraint(t*x[i] + (1-t)*x[i+1], t*y[i] + (1-t)*y[i+1]))
    return z3.And(c)

def step_length_constraint(x, y, maxd):
    c = []
    N = len(x)
    for i in range(N-1):
        c.append(z3.Implies(x[i+1] >  x[i], x[i+1]-x[i] <= maxd))
        c.append(z3.Implies(x[i+1] <= x[i], x[i]-x[i+1] <= maxd))
        c.append(z3.Implies(y[i+1] >  y[i], y[i+1]-y[i] <= maxd))
        c.append(z3.Implies(y[i+1] <= y[i], y[i]-y[i+1] <= maxd))
    return z3.And(c)

def abs_var(x):
    return z3.If(x >= 0, x, -x)

def trajectory_distance(x, y):
    d = 0
    N = len(x)
    for i in range(N-1):
        d += (x[i+1]-x[i])**2 + (y[i+1]-y[i])**2
    return d

def print_var(x, model):
    if isinstance(x, z3.ArithRef):
        val = eval(str(model.eval(x)))
        print(str(x) + ":" + str(val))

def get_var(x, model):
    if isinstance(x, z3.ArithRef):
        val = eval(str(model.eval(x)))
        return val
    else:
        return x

def free_single_linear_bool(x, y):
    belt1 = x-y > -1 and   x-y < -0.5 and  x<0 and   y>0 and   y<0.25
    belt2 = x+y > 0.5 and  x+y < 1 and     x>=0 and  y>0 and   y<0.25
    belt3 = x-y > 0.5 and  x-y < 1 and     x>=0 and  y<=0
    belt4 = x+y > -1 and   x+y < -0.5 and  x<0 and   y<=0
    broken_square1 = belt1 or  belt2 or  belt3 or  belt4
    
    belt1 = x-y > -2 and   x-y < -1.5 and  x<0 and   y>0
    belt2 = x+y > 1.5 and  x+y < 2 and     x>=0 and  y>0
    belt3 = x-y > 1.5 and  x-y < 2 and     x>=0 and  y<=0 and  y>-1.25
    belt4 = x+y > -2 and   x+y < -1.5 and  x<0 and   y<=0 and  y>-1.25
    broken_square2 = belt1 or  belt2 or  belt3 or  belt4
    
    obs = broken_square1 or broken_square2
    free = not obs
    return free

def free_single_easy_bool(x, y):
    belt1 = x < -0.5 and x > -1 and y < 1 and y > -1
    belt2 = x < 1    and x > 0.5 and y < 1 and y > -1
    belt3 = y < -0.5 and y > -1  and x < 1 and x > -1
    broken_square1 = belt1 or belt2 or belt3
    belt1 = x < -1.5 and x > -2 and  y < 2 and y > -2
    belt2 = x < 2 and    x > 1.5 and y < 2 and y > -2
    belt3 = y < 2 and y > 1.5 and  x < 2 and x > -2
    broken_square2 = belt1 or belt2 or belt3

    obs = broken_square1 or broken_square2
    free = not obs
    return free

